"""
图像生成任务模型
用于管理用户的AI图像生成请求，支持高并发和任务追踪
"""

from sqlalchemy import Column, String, DateTime, Integer, Text, Float, ForeignKey, CHAR
from sqlalchemy.orm import relationship
from datetime import datetime
import enum
import uuid

from .base import Base


class TaskStatus(str, enum.Enum):
    """任务状态枚举"""
    PENDING = "pending"           # 待处理
    PROCESSING = "processing"     # 处理中
    PROMPT_REFINING = "prompt_refining"  # 提示词提炼中
    COMFYUI_QUEUED = "comfyui_queued"    # ComfyUI队列中
    COMFYUI_PROCESSING = "comfyui_processing"  # ComfyUI处理中
    COMPLETED = "completed"       # 已完成
    FAILED = "failed"            # 失败
    CANCELLED = "cancelled"      # 已取消
    TIMEOUT = "timeout"          # 超时


class ImageGenerationTask(Base):
    """
    图像生成任务模型
    记录用户的图像生成请求和处理状态
    """
    __tablename__ = "image_generation_tasks"
    
    # 主键 - 使用CHAR(36)存储UUID字符串
    task_id = Column(CHAR(36), primary_key=True, default=lambda: str(uuid.uuid4()), nullable=False, comment="任务唯一标识符")
    
    # 外键关联
    user_uuid = Column(CHAR(36), ForeignKey("visitors.uuid"), nullable=False, comment="用户UUID")
    ai_person_id = Column(Integer, ForeignKey("ai_persons.id"), nullable=True, comment="关联的AI角色ID")
    
    # 任务基本信息
    status = Column(String, nullable=False, default=TaskStatus.PENDING.value, comment="任务状态")
    priority = Column(Integer, nullable=False, default=5, comment="任务优先级（1-10，数字越小优先级越高）")
    
    # 用户输入
    character_trigger = Column(Text, nullable=False, comment="角色触发词")
    user_prompt = Column(Text, nullable=False, comment="用户原始提示词")
    refined_prompt = Column(Text, nullable=True, comment="提炼后的英文提示词")
    
    # ComfyUI相关
    comfyui_prompt_id = Column(String, nullable=True, comment="ComfyUI任务ID")
    comfyui_workflow = Column(Text, nullable=True, comment="发送给ComfyUI的完整工作流JSON")
    
    # 图像生成配置（固定参数）
    image_width = Column(Integer, nullable=False, default=720, comment="图像宽度")
    image_height = Column(Integer, nullable=False, default=1280, comment="图像高度")
    sampling_steps = Column(Integer, nullable=False, default=25, comment="采样步数")
    
    # 结果信息
    image_filename = Column(String, nullable=True, comment="生成的图像文件名")
    image_url = Column(String, nullable=True, comment="图像访问URL")
    image_size_bytes = Column(Integer, nullable=True, comment="图像文件大小（字节）")
    
    # 处理时间统计
    created_at = Column(DateTime, nullable=False, default=datetime.utcnow, comment="任务创建时间")
    started_at = Column(DateTime, nullable=True, comment="任务开始处理时间")
    prompt_refined_at = Column(DateTime, nullable=True, comment="提示词提炼完成时间")
    comfyui_submitted_at = Column(DateTime, nullable=True, comment="提交到ComfyUI时间")
    completed_at = Column(DateTime, nullable=True, comment="任务完成时间")
    updated_at = Column(DateTime, nullable=False, default=datetime.utcnow, onupdate=datetime.utcnow, comment="最后更新时间")
    
    # 错误信息
    error_message = Column(Text, nullable=True, comment="错误信息")
    retry_count = Column(Integer, nullable=False, default=0, comment="重试次数")
    max_retries = Column(Integer, nullable=False, default=3, comment="最大重试次数")
    
    # 性能指标
    total_processing_time = Column(Float, nullable=True, comment="总处理时间（秒）")
    prompt_refining_time = Column(Float, nullable=True, comment="提示词提炼耗时（秒）")
    comfyui_processing_time = Column(Float, nullable=True, comment="ComfyUI处理耗时（秒）")
    
    # 客户端信息（用于调试和分析）
    client_ip = Column(String, nullable=True, comment="客户端IP地址")
    user_agent = Column(String, nullable=True, comment="用户代理字符串")
    
    def __repr__(self):
        return f"<ImageGenerationTask(task_id={self.task_id}, status={self.status}, user_uuid={self.user_uuid})>"
    
    def is_terminal_status(self) -> bool:
        """检查是否为终端状态（已完成、失败、取消）"""
        terminal_statuses = [
            TaskStatus.COMPLETED.value,
            TaskStatus.FAILED.value,
            TaskStatus.CANCELLED.value,
            TaskStatus.TIMEOUT.value
        ]
        return self.status in terminal_statuses
    
    def can_retry(self) -> bool:
        """检查是否可以重试"""
        return (
            self.status == TaskStatus.FAILED.value and 
            self.retry_count < self.max_retries
        )
    
    def calculate_processing_time(self) -> float:
        """计算总处理时间"""
        if self.started_at and self.completed_at:
            return (self.completed_at - self.started_at).total_seconds()
        return 0.0
    
    def get_progress_percentage(self) -> int:
        """获取任务进度百分比"""
        status_progress = {
            TaskStatus.PENDING.value: 0,
            TaskStatus.PROCESSING.value: 10,
            TaskStatus.PROMPT_REFINING.value: 20,
            TaskStatus.COMFYUI_QUEUED.value: 30,
            TaskStatus.COMFYUI_PROCESSING.value: 60,
            TaskStatus.COMPLETED.value: 100,
            TaskStatus.FAILED.value: 0,
            TaskStatus.CANCELLED.value: 0,
            TaskStatus.TIMEOUT.value: 0
        }
        return status_progress.get(self.status, 0)
    
    def update_status(self, new_status: str, error_message: str = None):
        """
        更新任务状态并记录时间戳
        
        Args:
            new_status: 新的任务状态
            error_message: 错误信息（可选）
        """
        self.status = new_status
        self.updated_at = datetime.utcnow()
        
        # 记录特定状态的时间戳
        if new_status == TaskStatus.PROCESSING.value and not self.started_at:
            self.started_at = datetime.utcnow()
        elif new_status == TaskStatus.PROMPT_REFINING.value:
            self.prompt_refined_at = datetime.utcnow()
        elif new_status == TaskStatus.COMFYUI_QUEUED.value:
            self.comfyui_submitted_at = datetime.utcnow()
        elif new_status in [TaskStatus.COMPLETED.value, TaskStatus.FAILED.value, TaskStatus.TIMEOUT.value]:
            self.completed_at = datetime.utcnow()
            # 计算总处理时间
            if self.started_at:
                self.total_processing_time = self.calculate_processing_time()
        
        # 记录错误信息
        if error_message:
            self.error_message = error_message
            
        # 失败时增加重试计数
        if new_status == TaskStatus.FAILED.value:
            self.retry_count += 1
    
    def to_dict(self) -> dict:
        """转换为字典格式，用于API响应"""
        return {
            "task_id": str(self.task_id),
            "status": self.status,
            "progress": self.get_progress_percentage(),
            "user_prompt": self.user_prompt,
            "refined_prompt": self.refined_prompt,
            "character_trigger": self.character_trigger,
            "image_url": self.image_url,
            "image_filename": self.image_filename,
            "created_at": self.created_at.isoformat() if self.created_at else None,
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
            "error_message": self.error_message,
            "retry_count": self.retry_count,
            "total_processing_time": self.total_processing_time,
            "estimated_time_remaining": self._estimate_remaining_time()
        }
    
    def _estimate_remaining_time(self) -> int:
        """估算剩余处理时间（秒）"""
        if self.status == TaskStatus.COMPLETED.value:
            return 0
        elif self.status in [TaskStatus.FAILED.value, TaskStatus.CANCELLED.value, TaskStatus.TIMEOUT.value]:
            return 0
        elif self.status == TaskStatus.PENDING.value:
            return 120  # 预估2分钟
        elif self.status in [TaskStatus.PROCESSING.value, TaskStatus.PROMPT_REFINING.value]:
            return 100  # 预估1分40秒
        elif self.status == TaskStatus.COMFYUI_QUEUED.value:
            return 80   # 预估1分20秒
        elif self.status == TaskStatus.COMFYUI_PROCESSING.value:
            return 40   # 预估40秒
        else:
            return 60   # 默认1分钟 